import dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate,FewShotChatMessagePromptTemplate
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()
# 1、构建示例模板与示例

example = [
    {"question":"帮我计算一下2+2等于多少","answer":"4"},
    {"question":"帮我计算一下2+3等于多少","answer":"5"},
    {"question":"帮我计算一下2+12等于多少","answer":"14"}
]

example_prompt = ChatPromptTemplate.from_messages([
    ("human","{question}"),
    ("ai","{answer}")
])

# 构建少量提示词模板
few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt = example_prompt,
    examples = example
)

# 3、构建最终提示词模板
prompt = ChatPromptTemplate.from_messages([
    ("system","你是一个可以计算复杂数学问题的聊天机器人，请严格按照示例样本输出，不要废话"),
    few_shot_prompt,
    ("human","{question}")
])

# 4、创建大语言模型与链
llm = ChatOpenAI(model="gpt-4o-mini",temperature=0)
chain =prompt | llm | StrOutputParser()

# 5、调用链获取结果
print(chain.invoke("帮我计算5*8等于多少"))